{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 可视化展示模型去雾效果\n",
    "## 第一步：设置环境变量"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ecd5972b561aad73"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: warning: setlocale: LC_ALL: cannot change locale (en_US.UTF-8)\r\n",
      "/bin/bash: warning: setlocale: LC_ALL: cannot change locale (en_US.UTF-8)\r\n",
      "/bin/bash: warning: setlocale: LC_ALL: cannot change locale (en_US.UTF-8)\r\n"
     ]
    }
   ],
   "source": [
    "import platform\n",
    "!BASICSR_JIT=True\n",
    "platform = platform.system()\n",
    "if platform == 'Linux':\n",
    "    !BASICSR_JIT=True\n",
    "else:\n",
    "    !set BASICSR_JIT=True\n",
    "!CUDA_VISIBLE_DEVICES=0\n",
    "\n"
   ],
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:56:07.763261Z",
     "start_time": "2024-03-01T14:56:07.410100Z"
    }
   },
   "id": "initial_id",
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 第二步：导入相关依赖"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7b426cdcb2ddc093"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from basicsr.archs.itb_arch import FusionRefine\n",
    "import torch\n",
    "import os\n",
    "import pyiqa\n",
    "import glob\n",
    "from tqdm import tqdm\n",
    "import torchvision\n",
    "from torchvision.transforms import ToTensor\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:47:12.331551Z",
     "start_time": "2024-03-01T14:47:12.329254Z"
    }
   },
   "id": "beeca1ed2dca5138",
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 第三步：定义变量"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8cf43fc7b109d94f"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "# 指定模型运算设备\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# 指定预训练模型存放位置\n",
    "pretrained_net_path = \"/mnt/e/DeepLearningCopies/2023/RIDCP/pretrained_models/ITB-Train-Best/DENSE-HAZE-2042-06374.pth\"\n",
    "# 指定待评估图片路径\n",
    "haze_img_path = \"/mnt/d/DeepLearning/dataset/Dense-Haze/hazy\"\n",
    "# 指定输出图像保存路径\n",
    "output_img_path = \"/mnt/e/DeepLearningCopies/2023/RIDCP/ohaze_results/DenseHaze\"\n",
    "# 指定图像最大分辨率\n",
    "# 分辨率过大容易爆显存，超过最大分辨率的将会降采样后交由模型处理\n",
    "max_size = 1000 * 1000"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:56:09.923181Z",
     "start_time": "2024-03-01T14:56:09.920316Z"
    }
   },
   "id": "e92940ec4e158eb0",
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 第四步：构建模型"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "942eed4917202f45"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: warning: setlocale: LC_ALL: cannot change locale (en_US.UTF-8)\r\n",
      "{'LQ_stage': True, 'use_weight': False, 'weight_alpha': -21.25}\n",
      "{}\n"
     ]
    },
    {
     "data": {
      "text/plain": "FusionRefine(\n  (feature_extract): RIDCPNew(\n    (vq_encoder): VQEncoder(\n      (in_conv): Conv2d(3, 64, kernel_size=(4, 4), stride=(1, 1), padding=(1, 1))\n      (blocks): ModuleList(\n        (0): Sequential(\n          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n          (1): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (2): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n        )\n        (1): Sequential(\n          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n          (1): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 256, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 256, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (2): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 256, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 256, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n        )\n      )\n    )\n    (ridcp_encoder): SwinLayers(\n      (swin_blks): ModuleList(\n        (0-3): 4 x RSTB(\n          (residual_group): BasicLayer(\n            dim=256, input_resolution=(32, 32), depth=6\n            (blocks): ModuleList(\n              (0): SwinTransformerBlock(\n                dim=256, input_resolution=(32, 32), num_heads=8, window_size=8, shift_size=0, mlp_ratio=4.0\n                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (attn): WindowAttention(\n                  dim=256, window_size=(8, 8), num_heads=8\n                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n                  (attn_drop): Dropout(p=0.0, inplace=False)\n                  (proj): Linear(in_features=256, out_features=256, bias=True)\n                  (proj_drop): Dropout(p=0.0, inplace=False)\n                  (softmax): Softmax(dim=-1)\n                )\n                (drop_path): Identity()\n                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (mlp): Mlp(\n                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n                  (act): GELU(approximate='none')\n                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n                  (drop): Dropout(p=0.0, inplace=False)\n                )\n              )\n              (1): SwinTransformerBlock(\n                dim=256, input_resolution=(32, 32), num_heads=8, window_size=8, shift_size=4, mlp_ratio=4.0\n                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (attn): WindowAttention(\n                  dim=256, window_size=(8, 8), num_heads=8\n                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n                  (attn_drop): Dropout(p=0.0, inplace=False)\n                  (proj): Linear(in_features=256, out_features=256, bias=True)\n                  (proj_drop): Dropout(p=0.0, inplace=False)\n                  (softmax): Softmax(dim=-1)\n                )\n                (drop_path): Identity()\n                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (mlp): Mlp(\n                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n                  (act): GELU(approximate='none')\n                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n                  (drop): Dropout(p=0.0, inplace=False)\n                )\n              )\n              (2): SwinTransformerBlock(\n                dim=256, input_resolution=(32, 32), num_heads=8, window_size=8, shift_size=0, mlp_ratio=4.0\n                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (attn): WindowAttention(\n                  dim=256, window_size=(8, 8), num_heads=8\n                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n                  (attn_drop): Dropout(p=0.0, inplace=False)\n                  (proj): Linear(in_features=256, out_features=256, bias=True)\n                  (proj_drop): Dropout(p=0.0, inplace=False)\n                  (softmax): Softmax(dim=-1)\n                )\n                (drop_path): Identity()\n                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (mlp): Mlp(\n                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n                  (act): GELU(approximate='none')\n                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n                  (drop): Dropout(p=0.0, inplace=False)\n                )\n              )\n              (3): SwinTransformerBlock(\n                dim=256, input_resolution=(32, 32), num_heads=8, window_size=8, shift_size=4, mlp_ratio=4.0\n                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (attn): WindowAttention(\n                  dim=256, window_size=(8, 8), num_heads=8\n                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n                  (attn_drop): Dropout(p=0.0, inplace=False)\n                  (proj): Linear(in_features=256, out_features=256, bias=True)\n                  (proj_drop): Dropout(p=0.0, inplace=False)\n                  (softmax): Softmax(dim=-1)\n                )\n                (drop_path): Identity()\n                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (mlp): Mlp(\n                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n                  (act): GELU(approximate='none')\n                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n                  (drop): Dropout(p=0.0, inplace=False)\n                )\n              )\n              (4): SwinTransformerBlock(\n                dim=256, input_resolution=(32, 32), num_heads=8, window_size=8, shift_size=0, mlp_ratio=4.0\n                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (attn): WindowAttention(\n                  dim=256, window_size=(8, 8), num_heads=8\n                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n                  (attn_drop): Dropout(p=0.0, inplace=False)\n                  (proj): Linear(in_features=256, out_features=256, bias=True)\n                  (proj_drop): Dropout(p=0.0, inplace=False)\n                  (softmax): Softmax(dim=-1)\n                )\n                (drop_path): Identity()\n                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (mlp): Mlp(\n                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n                  (act): GELU(approximate='none')\n                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n                  (drop): Dropout(p=0.0, inplace=False)\n                )\n              )\n              (5): SwinTransformerBlock(\n                dim=256, input_resolution=(32, 32), num_heads=8, window_size=8, shift_size=4, mlp_ratio=4.0\n                (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (attn): WindowAttention(\n                  dim=256, window_size=(8, 8), num_heads=8\n                  (qkv): Linear(in_features=256, out_features=768, bias=True)\n                  (attn_drop): Dropout(p=0.0, inplace=False)\n                  (proj): Linear(in_features=256, out_features=256, bias=True)\n                  (proj_drop): Dropout(p=0.0, inplace=False)\n                  (softmax): Softmax(dim=-1)\n                )\n                (drop_path): Identity()\n                (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n                (mlp): Mlp(\n                  (fc1): Linear(in_features=256, out_features=1024, bias=True)\n                  (act): GELU(approximate='none')\n                  (fc2): Linear(in_features=1024, out_features=256, bias=True)\n                  (drop): Dropout(p=0.0, inplace=False)\n                )\n              )\n            )\n          )\n          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (patch_embed): PatchEmbed()\n          (patch_unembed): PatchUnEmbed()\n        )\n      )\n    )\n    (ridcp_decoder): RIDCPDecoder(\n      (upsampler): ModuleList(\n        (0): Sequential(\n          (0): Upsample(scale_factor=2.0, mode='nearest')\n          (1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (2): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (3): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (4): DehazeBlock(\n            (conv1): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n            (act1): ReLU(inplace=True)\n            (conv2): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n            (calayer): CALayer(\n              (avg_pool): AdaptiveAvgPool2d(output_size=1)\n              (ca): Sequential(\n                (0): Conv2d(128, 16, kernel_size=(1, 1), stride=(1, 1))\n                (1): ReLU(inplace=True)\n                (2): Conv2d(16, 128, kernel_size=(1, 1), stride=(1, 1))\n                (3): Sigmoid()\n              )\n            )\n            (palayer): PALayer(\n              (pa): Sequential(\n                (0): Conv2d(128, 16, kernel_size=(1, 1), stride=(1, 1))\n                (1): ReLU(inplace=True)\n                (2): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1))\n                (3): Sigmoid()\n              )\n            )\n          )\n        )\n        (1): Sequential(\n          (0): Upsample(scale_factor=2.0, mode='nearest')\n          (1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (2): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (3): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (4): DehazeBlock(\n            (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n            (act1): ReLU(inplace=True)\n            (conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n            (calayer): CALayer(\n              (avg_pool): AdaptiveAvgPool2d(output_size=1)\n              (ca): Sequential(\n                (0): Conv2d(64, 8, kernel_size=(1, 1), stride=(1, 1))\n                (1): ReLU(inplace=True)\n                (2): Conv2d(8, 64, kernel_size=(1, 1), stride=(1, 1))\n                (3): Sigmoid()\n              )\n            )\n            (palayer): PALayer(\n              (pa): Sequential(\n                (0): Conv2d(64, 8, kernel_size=(1, 1), stride=(1, 1))\n                (1): ReLU(inplace=True)\n                (2): Conv2d(8, 1, kernel_size=(1, 1), stride=(1, 1))\n                (3): Sigmoid()\n              )\n            )\n          )\n        )\n      )\n      (warp): ModuleList(\n        (0): WarpBlock(\n          (offset): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (dcn): DCNv2Pack(\n            (conv_offset): Conv2d(128, 108, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          )\n        )\n        (1): WarpBlock(\n          (offset): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (dcn): DCNv2Pack(\n            (conv_offset): Conv2d(64, 108, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          )\n        )\n      )\n    )\n    (vq_decoder_group): ModuleList(\n      (0): DecoderBlock(\n        (block): Sequential(\n          (0): Upsample(scale_factor=2.0, mode='nearest')\n          (1): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (2): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (3): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 128, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n        )\n      )\n      (1): DecoderBlock(\n        (block): Sequential(\n          (0): Upsample(scale_factor=2.0, mode='nearest')\n          (1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n          (2): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n          (3): ResBlock(\n            (conv): Sequential(\n              (0): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (1): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): NormLayer(\n                (norm): GroupNorm(32, 64, eps=1e-06, affine=True)\n              )\n              (4): ActLayer(\n                (func): SiLU(inplace=True)\n              )\n              (5): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n            )\n          )\n        )\n      )\n    )\n    (out_conv): Conv2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (enhancer): Enhancer(\n      (relu): LeakyReLU(negative_slope=0.2, inplace=True)\n      (tanh): Tanh()\n      (refine1): Conv2d(3, 20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n      (refine2): Conv2d(20, 20, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n      (conv1010): Conv2d(20, 1, kernel_size=(1, 1), stride=(1, 1))\n      (conv1020): Conv2d(20, 1, kernel_size=(1, 1), stride=(1, 1))\n      (conv1030): Conv2d(20, 1, kernel_size=(1, 1), stride=(1, 1))\n      (conv1040): Conv2d(20, 1, kernel_size=(1, 1), stride=(1, 1))\n      (refine3): Conv2d(24, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n      (batch1): InstanceNorm2d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False)\n    )\n    (residual_conv): Conv2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (quantizer): VectorQuantizer(\n      (embedding): Embedding(1024, 512)\n    )\n    (before_quant): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1))\n    (after_quant): CombineQuantBlock(\n      (conv): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    )\n  )\n  (pre_trained_rcan): RCAN(\n    (head): Sequential(\n      (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    )\n    (body): Sequential(\n      (0): ResidualGroup(\n        (body): Sequential(\n          (0): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (1): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (2): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (3): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (4): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (5): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (6): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (7): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (8): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (9): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (10): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n        )\n      )\n      (1): ResidualGroup(\n        (body): Sequential(\n          (0): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (1): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (2): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (3): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (4): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (5): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (6): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (7): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (8): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (9): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (10): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n        )\n      )\n      (2): ResidualGroup(\n        (body): Sequential(\n          (0): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (1): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (2): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (3): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (4): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (5): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (6): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (7): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (8): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (9): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (10): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n        )\n      )\n      (3): ResidualGroup(\n        (body): Sequential(\n          (0): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (1): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (2): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (3): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (4): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (5): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (6): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (7): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (8): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (9): RCAB(\n            (body): Sequential(\n              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (1): ReLU(inplace=True)\n              (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n              (3): RC_CALayer(\n                (avg_pool): AdaptiveAvgPool2d(output_size=1)\n                (conv_du): Sequential(\n                  (0): Conv2d(32, 4, kernel_size=(1, 1), stride=(1, 1))\n                  (1): ReLU()\n                  (2): Conv2d(4, 32, kernel_size=(1, 1), stride=(1, 1))\n                  (3): Sigmoid()\n                )\n              )\n            )\n          )\n          (10): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n        )\n      )\n      (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    )\n    (tail): Sequential(\n      (0): Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    )\n  )\n  (tail): Sequential(\n    (0): ReflectionPad2d((3, 3, 3, 3))\n    (1): Conv2d(35, 3, kernel_size=(7, 7), stride=(1, 1))\n    (2): Tanh()\n  )\n)"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "!BASICSR_JIT=True\n",
    "# 构建模型，加载预训练权重\n",
    "opt = {\n",
    "    \"LQ_stage\": True,\n",
    "    \"use_weight\": False,\n",
    "    \"weight_alpha\": -21.25\n",
    "}\n",
    "sr_model = FusionRefine(opt=opt).to(device)\n",
    "sr_model.load_state_dict(torch.load(pretrained_net_path)['params'], strict=False)\n",
    "sr_model.eval()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:56:11.934270Z",
     "start_time": "2024-03-01T14:56:11.132478Z"
    }
   },
   "id": "2f190f83483e49c3",
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 第五步：处理有雾图像"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "63f112b5e9f7284b"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "psnr = pyiqa.create_metric(\"psnr\", device=device)\n",
    "ssim = pyiqa.create_metric(\"ssim\", device=device)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:56:14.574714Z",
     "start_time": "2024-03-01T14:56:14.571832Z"
    }
   },
   "id": "4a8bbb95505d0e6d",
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "if os.path.isfile(haze_img_path):\n",
    "    paths = [haze_img_path]\n",
    "else:\n",
    "    paths = sorted(glob.glob(os.path.join(haze_img_path, '*')))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:56:16.493460Z",
     "start_time": "2024-03-01T14:56:16.486431Z"
    }
   },
   "id": "4afe03397a2f571e",
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "处理图像 01_hazy.png 中:   0%|          | 0/55 [00:39<?, ?image/s]\n",
      "\n",
      "处理图像 01_hazy.png 中:   0%|          | 0/55 [00:00<?, ?image/s]\u001B[A"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'deform_conv_ext' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[12], line 20\u001B[0m\n\u001B[1;32m     18\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m     19\u001B[0m     input_img \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mnn\u001B[38;5;241m.\u001B[39mUpsamplingBilinear2d((h\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m2\u001B[39m, w\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m/\u001B[39m\u001B[38;5;241m2\u001B[39m))(input_img)\n\u001B[0;32m---> 20\u001B[0m     output, _ \u001B[38;5;241m=\u001B[39m \u001B[43msr_model\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mtest\u001B[49m\u001B[43m(\u001B[49m\u001B[43minput_img\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     21\u001B[0m     output \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mnn\u001B[38;5;241m.\u001B[39mUpsamplingBilinear2d((h, w))(output)\n\u001B[1;32m     23\u001B[0m torchvision\u001B[38;5;241m.\u001B[39mutils\u001B[38;5;241m.\u001B[39msave_image(output, save_path)\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/utils/_contextlib.py:115\u001B[0m, in \u001B[0;36mcontext_decorator.<locals>.decorate_context\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    112\u001B[0m \u001B[38;5;129m@functools\u001B[39m\u001B[38;5;241m.\u001B[39mwraps(func)\n\u001B[1;32m    113\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mdecorate_context\u001B[39m(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs):\n\u001B[1;32m    114\u001B[0m     \u001B[38;5;28;01mwith\u001B[39;00m ctx_factory():\n\u001B[0;32m--> 115\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mfunc\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/itb_arch.py:46\u001B[0m, in \u001B[0;36mFusionRefine.test\u001B[0;34m(self, inputs)\u001B[0m\n\u001B[1;32m     44\u001B[0m \u001B[38;5;129m@torch\u001B[39m\u001B[38;5;241m.\u001B[39mno_grad()\n\u001B[1;32m     45\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mtest\u001B[39m(\u001B[38;5;28mself\u001B[39m, inputs):\n\u001B[0;32m---> 46\u001B[0m     x, index \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mfeature_extract\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mtest\u001B[49m\u001B[43m(\u001B[49m\u001B[43minputs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     47\u001B[0m     y \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mpre_trained_rcan(inputs)\n\u001B[1;32m     48\u001B[0m     out \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mcat([x, y], \u001B[38;5;241m1\u001B[39m)\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/utils/_contextlib.py:115\u001B[0m, in \u001B[0;36mcontext_decorator.<locals>.decorate_context\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    112\u001B[0m \u001B[38;5;129m@functools\u001B[39m\u001B[38;5;241m.\u001B[39mwraps(func)\n\u001B[1;32m    113\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mdecorate_context\u001B[39m(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs):\n\u001B[1;32m    114\u001B[0m     \u001B[38;5;28;01mwith\u001B[39;00m ctx_factory():\n\u001B[0;32m--> 115\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mfunc\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/ridcp_new_arch.py:210\u001B[0m, in \u001B[0;36mRIDCPNew.test\u001B[0;34m(self, inputs)\u001B[0m\n\u001B[1;32m    207\u001B[0m inputs \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mcat([inputs, torch\u001B[38;5;241m.\u001B[39mflip(inputs, [\u001B[38;5;241m2\u001B[39m])], \u001B[38;5;241m2\u001B[39m)[:, :, :h_old \u001B[38;5;241m+\u001B[39m h_pad, :]\n\u001B[1;32m    208\u001B[0m inputs \u001B[38;5;241m=\u001B[39m torch\u001B[38;5;241m.\u001B[39mcat([inputs, torch\u001B[38;5;241m.\u001B[39mflip(inputs, [\u001B[38;5;241m3\u001B[39m])], \u001B[38;5;241m3\u001B[39m)[:, :, :, :w_old \u001B[38;5;241m+\u001B[39m w_pad]\n\u001B[0;32m--> 210\u001B[0m output_vq, output, _, _, after_quant, index \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mencode_and_decode\u001B[49m\u001B[43m(\u001B[49m\u001B[43minputs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\u001B[43m)\u001B[49m\n\u001B[1;32m    212\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m output \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[1;32m    213\u001B[0m     output \u001B[38;5;241m=\u001B[39m output[\u001B[38;5;241m.\u001B[39m\u001B[38;5;241m.\u001B[39m\u001B[38;5;241m.\u001B[39m, :h_old, :w_old]\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/ridcp_new_arch.py:125\u001B[0m, in \u001B[0;36mRIDCPNew.encode_and_decode\u001B[0;34m(self, inputs, gt_indices)\u001B[0m\n\u001B[1;32m    123\u001B[0m         residual_feature \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mridcp_decoder(enc_feats, code_decoder_output)\n\u001B[1;32m    124\u001B[0m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m--> 125\u001B[0m         residual_feature \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mridcp_decoder\u001B[49m\u001B[43m(\u001B[49m\u001B[43menc_feats\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdetach\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcode_decoder_output\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    126\u001B[0m     out_img_residual \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mresidual_conv(residual_feature)\n\u001B[1;32m    127\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001B[0m, in \u001B[0;36mModule._wrapped_call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1516\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_compiled_call_impl(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m   1517\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1518\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_impl\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1522\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[1;32m   1523\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[1;32m   1524\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[1;32m   1525\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[1;32m   1526\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[0;32m-> 1527\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mforward_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1529\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m   1530\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/ridcp/decoder.py:132\u001B[0m, in \u001B[0;36mRIDCPDecoder.forward\u001B[0;34m(self, x, code_decoder_output)\u001B[0m\n\u001B[1;32m    130\u001B[0m x \u001B[38;5;241m=\u001B[39m m(x)\n\u001B[1;32m    131\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39muse_warp:\n\u001B[0;32m--> 132\u001B[0m     x_vq \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mwarp\u001B[49m\u001B[43m[\u001B[49m\u001B[43midx\u001B[49m\u001B[43m]\u001B[49m\u001B[43m(\u001B[49m\u001B[43mcode_decoder_output\u001B[49m\u001B[43m[\u001B[49m\u001B[43midx\u001B[49m\u001B[43m]\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mx\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    133\u001B[0m     x \u001B[38;5;241m=\u001B[39m x \u001B[38;5;241m+\u001B[39m x_vq \u001B[38;5;241m*\u001B[39m (x\u001B[38;5;241m.\u001B[39mmean() \u001B[38;5;241m/\u001B[39m x_vq\u001B[38;5;241m.\u001B[39mmean())\n\u001B[1;32m    134\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001B[0m, in \u001B[0;36mModule._wrapped_call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1516\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_compiled_call_impl(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m   1517\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1518\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_impl\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1522\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[1;32m   1523\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[1;32m   1524\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[1;32m   1525\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[1;32m   1526\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[0;32m-> 1527\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mforward_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1529\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m   1530\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/module/util_block.py:137\u001B[0m, in \u001B[0;36mWarpBlock.forward\u001B[0;34m(self, x_vq, x_residual)\u001B[0m\n\u001B[1;32m    135\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mforward\u001B[39m(\u001B[38;5;28mself\u001B[39m, x_vq, x_residual):\n\u001B[1;32m    136\u001B[0m     x_residual \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39moffset(torch\u001B[38;5;241m.\u001B[39mcat([x_vq, x_residual], dim\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m1\u001B[39m))\n\u001B[0;32m--> 137\u001B[0m     feat_after_warp \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdcn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mx_vq\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mx_residual\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    139\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m feat_after_warp\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001B[0m, in \u001B[0;36mModule._wrapped_call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1516\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_compiled_call_impl(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m   1517\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1518\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_impl\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1522\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[1;32m   1523\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[1;32m   1524\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[1;32m   1525\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[1;32m   1526\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[0;32m-> 1527\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mforward_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1529\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m   1530\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/module/util_block.py:125\u001B[0m, in \u001B[0;36mDCNv2Pack.forward\u001B[0;34m(self, x, feat)\u001B[0m\n\u001B[1;32m    122\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m offset_absmean \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m50\u001B[39m:\n\u001B[1;32m    123\u001B[0m     \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;124mf\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mOffset abs mean is \u001B[39m\u001B[38;5;132;01m{\u001B[39;00moffset_absmean\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m, larger than 50.\u001B[39m\u001B[38;5;124m'\u001B[39m)\n\u001B[0;32m--> 125\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mmodulated_deform_conv\u001B[49m\u001B[43m(\u001B[49m\u001B[43mx\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43moffset\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mmask\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mweight\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbias\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstride\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mpadding\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    126\u001B[0m \u001B[43m                             \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdilation\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mgroups\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdeformable_groups\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/.local/lib/python3.11/site-packages/torch/autograd/function.py:539\u001B[0m, in \u001B[0;36mFunction.apply\u001B[0;34m(cls, *args, **kwargs)\u001B[0m\n\u001B[1;32m    536\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m torch\u001B[38;5;241m.\u001B[39m_C\u001B[38;5;241m.\u001B[39m_are_functorch_transforms_active():\n\u001B[1;32m    537\u001B[0m     \u001B[38;5;66;03m# See NOTE: [functorch vjp and autograd interaction]\u001B[39;00m\n\u001B[1;32m    538\u001B[0m     args \u001B[38;5;241m=\u001B[39m _functorch\u001B[38;5;241m.\u001B[39mutils\u001B[38;5;241m.\u001B[39munwrap_dead_wrappers(args)\n\u001B[0;32m--> 539\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43msuper\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mapply\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m    541\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mcls\u001B[39m\u001B[38;5;241m.\u001B[39msetup_context \u001B[38;5;241m==\u001B[39m _SingleLevelFunction\u001B[38;5;241m.\u001B[39msetup_context:\n\u001B[1;32m    542\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mRuntimeError\u001B[39;00m(\n\u001B[1;32m    543\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mIn order to use an autograd.Function with functorch transforms \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    544\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m(vmap, grad, jvp, jacrev, ...), it must override the setup_context \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    545\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mstaticmethod. For more details, please see \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    546\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mhttps://pytorch.org/docs/master/notes/extending.func.html\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    547\u001B[0m     )\n",
      "File \u001B[0;32m/mnt/e/DeepLearningCopies/2023/RIDCP/basicsr/archs/dcn/deform_conv.py:151\u001B[0m, in \u001B[0;36mModulatedDeformConvFunction.forward\u001B[0;34m(ctx, input, offset, mask, weight, bias, stride, padding, dilation, groups, deformable_groups)\u001B[0m\n\u001B[1;32m    149\u001B[0m output \u001B[38;5;241m=\u001B[39m \u001B[38;5;28minput\u001B[39m\u001B[38;5;241m.\u001B[39mnew_empty(ModulatedDeformConvFunction\u001B[38;5;241m.\u001B[39m_infer_shape(ctx, \u001B[38;5;28minput\u001B[39m, weight))\n\u001B[1;32m    150\u001B[0m ctx\u001B[38;5;241m.\u001B[39m_bufs \u001B[38;5;241m=\u001B[39m [\u001B[38;5;28minput\u001B[39m\u001B[38;5;241m.\u001B[39mnew_empty(\u001B[38;5;241m0\u001B[39m), \u001B[38;5;28minput\u001B[39m\u001B[38;5;241m.\u001B[39mnew_empty(\u001B[38;5;241m0\u001B[39m)]\n\u001B[0;32m--> 151\u001B[0m \u001B[43mdeform_conv_ext\u001B[49m\u001B[38;5;241m.\u001B[39mmodulated_deform_conv_forward(\u001B[38;5;28minput\u001B[39m, weight, bias, ctx\u001B[38;5;241m.\u001B[39m_bufs[\u001B[38;5;241m0\u001B[39m], offset, mask, output,\n\u001B[1;32m    152\u001B[0m                                               ctx\u001B[38;5;241m.\u001B[39m_bufs[\u001B[38;5;241m1\u001B[39m], weight\u001B[38;5;241m.\u001B[39mshape[\u001B[38;5;241m2\u001B[39m], weight\u001B[38;5;241m.\u001B[39mshape[\u001B[38;5;241m3\u001B[39m], ctx\u001B[38;5;241m.\u001B[39mstride,\n\u001B[1;32m    153\u001B[0m                                               ctx\u001B[38;5;241m.\u001B[39mstride, ctx\u001B[38;5;241m.\u001B[39mpadding, ctx\u001B[38;5;241m.\u001B[39mpadding, ctx\u001B[38;5;241m.\u001B[39mdilation, ctx\u001B[38;5;241m.\u001B[39mdilation,\n\u001B[1;32m    154\u001B[0m                                               ctx\u001B[38;5;241m.\u001B[39mgroups, ctx\u001B[38;5;241m.\u001B[39mdeformable_groups, ctx\u001B[38;5;241m.\u001B[39mwith_bias)\n\u001B[1;32m    155\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m output\n",
      "\u001B[0;31mNameError\u001B[0m: name 'deform_conv_ext' is not defined"
     ]
    }
   ],
   "source": [
    "images = []\n",
    "images_metric = []\n",
    "count = 0\n",
    "pbar = tqdm(total=len(paths), unit='image')\n",
    "for idx, path in enumerate(paths):\n",
    "    img_name = os.path.basename(path)\n",
    "    save_path = os.path.join(output_img_path, f'{img_name}')\n",
    "    pbar.set_description(f'处理图像 {img_name} 中')\n",
    "\n",
    "    input_img = ToTensor()(Image.open(path).convert('RGB')).to(device)[None, ::]\n",
    "    h, w = input_img.shape[2:]\n",
    "    if h * w < max_size:\n",
    "        output, _ = sr_model.test(input_img)\n",
    "    elif h * w > max_size * 2:\n",
    "        input_img = torch.nn.UpsamplingBilinear2d((h//3, w//3))(input_img)\n",
    "        output = sr_model.test_tile(input_img, tile_size=960, tile_pad=64)\n",
    "        output = torch.nn.UpsamplingBilinear2d((h, w))(output)\n",
    "    else:\n",
    "        input_img = torch.nn.UpsamplingBilinear2d((h//2, w//2))(input_img)\n",
    "        output, _ = sr_model.test(input_img)\n",
    "        output = torch.nn.UpsamplingBilinear2d((h, w))(output)\n",
    "        \n",
    "    torchvision.utils.save_image(output, save_path)\n",
    "    \n",
    "    psnr_hl = psnr(input_img, output).item()\n",
    "    ssim_hl = ssim(input_img, output).item()\n",
    "\n",
    "    clear_out = input_img.squeeze().permute(1, 2, 0)\n",
    "    hazy_out = output.squeeze().permute(1, 2, 0)\n",
    "    images_metric.append({\n",
    "        \"Name\": img_name,\n",
    "        \"PSNR\": psnr_hl,\n",
    "        \"SSIM\": ssim_hl,\n",
    "    })\n",
    "    images.append(clear_out)\n",
    "    images.append(hazy_out)\n",
    "    count += 1\n",
    "    pbar.update(1)\n",
    "pbar.close()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-01T14:56:18.720487Z",
     "start_time": "2024-03-01T14:56:17.963379Z"
    }
   },
   "id": "f29f1ca95e1266eb",
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 第六步：可视化模型去雾效果"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "80d2de3cf0cdda28"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "for i in range(count):\n",
    "    fig, axs = plt.subplots(1, 2)\n",
    "    for j, ax in enumerate(axs.flat):\n",
    "        print(images[i * 2 + j])\n",
    "        ax.imshow(images[i * 2 + j])\n",
    "        ax.axis('off')\n",
    "    plt.suptitle(\"Name: {} SSIM: {:.2f} PSNR: {:.2f}\".format(\n",
    "        images_metric[i][\"Name\"],\n",
    "        images_metric[i][\"SSIM\"],\n",
    "        images_metric[i][\"PSNR\"],\n",
    "    ))\n",
    "    plt.subplots_adjust(wspace=0, hspace=0, top=0.7)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8e84246ec5567ec8"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
